"""
多分类问题
Softmax 和 CrossEntropyLoss
"""
from torchvision import transforms  # 对图像数据进行原始处理
from torchvision import datasets
from torch.utils.data import DataLoader
import torch
import torch.optim as optim  # 优化器

# 数据处理
batch_size = 64
transform = transforms.Compose([  # Compose表示将[]中的操作按顺序一一进行
    # 将图像数据处理成张量，c*w*h的格式，c-通道数，w-宽度，h-高度
    # 并且将图像0-255的数据压缩到0-1之间
    transforms.ToTensor(),
    # 归一化(self, mean(均值), std(标准差), inplace=False)
    # 归一化方程：norm=(origin-mean)/std
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='../dataset/mnist/',
                               train=True,
                               download=True,
                               transform=transform)
train_loader = DataLoader(train_dataset,
                          shuffle=True,
                          batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/',
                              train=True,
                              download=False,
                              transform=transform)
test_loader = DataLoader(test_dataset,
                         shuffle=False,
                         batch_size=batch_size)


# 模型
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.l1 = torch.nn.Linear(784, 512)
        self.l2 = torch.nn.Linear(512, 256)
        self.l3 = torch.nn.Linear(256, 128)
        self.l4 = torch.nn.Linear(128, 64)
        self.l5 = torch.nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 784)  # 将输入数据分为n个维度为784的数据单元，-1表示会自动计算n的值
        x = torch.relu(self.l1(x))  # 使用relu激活函数
        x = torch.relu(self.l2(x))
        x = torch.relu(self.l3(x))
        x = torch.relu(self.l4(x))
        return self.l5(x)  # 最后一层不做激活操作，因为最后用softmax进行转换，在交叉熵损失中进行


model = Net()

# 损失函数,使用交叉熵损失
# 交叉熵损失=Softmax + loss(loss = -Y * log(Ypred))
criterion = torch.nn.CrossEntropyLoss()
# 优化器,使用随机梯度下降
# momentum表示冲量，默认值为0
# 每次更新x时，加上上次对x的更新*momentum，即v=-dx*lr+momentum*v
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)


def train(epoch):
    running_loss = 0.0
    # 训练集中每个batch进行一次训练
    for batch_index, data in enumerate(train_loader, 0):  # 0表示下标从0开始
        inputs, target = data  # 数据读取
        # print('inputs:',inputs,'  target:',target)
        optimizer.zero_grad()  # 初始化梯度

        # forward + backward + update
        outputs = model(inputs)
        loss = criterion(outputs, target)
        # print(loss, ' type:', type(loss))
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if batch_index % 300 == 299:
            print('[第%d轮,%5d个训练数据] loss: %.3f' % (epoch + 1, batch_index + 1, running_loss / 300))
            running_loss = 0.0


def test():
    correct = 0  # 预测正确的数量
    total = 0  # 总预测数
    with torch.no_grad():  # torch.no.grad()表示不计算梯度
        for data in test_loader:
            images, labels = data
            outputs = model(images)  # 这里的到的outputs对每个图片的预测是一个1*10的向量，最大值表示预测的分类
            # dim表示沿着第一个维度(即行维度)取最大值 （0表示列维度）
            # torch.max的返回值为(值，下标)
            _, predicted = torch.max(outputs.data, dim=1)
            total += labels.size(0)  # 总样本数
            correct += (predicted == labels).sum().item()  # 分类正确计数

    print('测试集上的准确率：%d %%' % (100 * correct / total))


if __name__ == '__main__':
    for epoch in range(10):
        train(epoch)
        test()
